
#include <benchmark/benchmark.h>
#include <err.h>
#include <getopt.h>
#include <inttypes.h>
#include <math.h>
#include <sys/resource.h>
#include <sys/stat.h>

#include <map>
#include <mutex>
#include <sstream>
#include <string>
#include <utility>
#include <vector>

//#include <tinyxml2.h>
#include "util.h"

#define STRCONNECT(x) #x

static const std::vector<int> commonArgs{
  8,
  16,
  32,
  64,
  512,
  1 * K,
  8 * K,
  16 * K,
  32 * K,
  64 * K,
  128 * K,
};

static const std::vector<int> smallArgs{
  // Increment by 1
  1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
  // Increment by 8
  24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144,
  // Increment by 16
  160, 176, 192, 208, 224, 240, 256,
};

static const std::vector<int> mediumArgs{
  512,
  1 * K,
  8 * K,
  16 * K,
  32 * K,
  64 * K,
  128 * K,
};

static const std::vector<int> largeArgs{
  256 * K,
  512 * K,
  1024 * K,
  2048 * K,
};

static std::map<std::string, const std::vector<int> &> argsSize{
  { "SMALL",  smallArgs },
  { "MEDIUM", mediumArgs },
  { "LARGE",  largeArgs },
};

std::map<std::string, std::pair<benchmark_func, std::string>> g_str_func_map;
std::mutex g_benchmark_map_lock;

static struct option g_long_options[] =
{
  {"musl_cpu", required_argument, nullptr, 'c'},
//  {"musl_xml", required_argument, nullptr, 'x'},
  {"musl_iterations", required_argument, nullptr, 'i'},
  {"musl_extra", required_argument, nullptr, 'a'},
  {"help", no_argument, nullptr, 'h'},
  {nullptr, 0, nullptr, 0},
};

typedef std::vector<std::vector<int64_t>> args_vector;

void Usage() {
  printf("Usage:\n");
  printf("musl_benchmarks [--musl_cpu=<cpu_to_isolate>]\n");
//  printf("                  [--musl_xml=<path_to_xml>]\n");
  printf("                  [--musl_iterations=<num_iter>]\n");
  printf("                  [--musl_extra=\"<fn_name> <arg1> <arg2> ...\"]\n");
  printf("                  [<original benchmark flags>]\n");
  printf("benchmark flags:\n");

  int argc = 2;
  char argv0[] = "musl_benchmarks";
  char argv1[] = "--help";
  char* argv[3] {argv0, argv1, nullptr};
  benchmark::Initialize(&argc, argv);
  exit(1);
}

void SanitizeOpts(int argc, char** argv, std::vector<char*>* argv_for_origin) {
  (*argv_for_origin)[0] = argv[0];
  for (int i = 1; i < argc; ++i) {
    char* optarg = argv[i];
    size_t opt_idx = 0;

    while (g_long_options[opt_idx].name && strncmp(g_long_options[opt_idx].name, optarg + 2,
                                                   strlen(g_long_options[opt_idx].name))) {
      ++opt_idx;
    }

    if (!g_long_options[opt_idx].name) {
      argv_for_origin->push_back(optarg);
    } else if ((g_long_options[opt_idx].has_arg == required_argument) && !strchr(optarg, '=')) {
      i++;
    }
  }
  argv_for_origin->push_back(nullptr);
}

bench_opts_t ParseOpts(int argc, char** argv) {
  bench_opts_t opts;
  int opt;
  int option_index = 0;

  // To make this parser handle the benchmark options silently:
  extern int opterr;
  opterr = 0;

  while ((opt = getopt_long(argc, argv, "c:x:i:a:h", g_long_options, &option_index)) != -1) {
    switch (opt) {
      case 'c':
        if (*optarg) {
          char* null_check;
          opts.cpu_num = strtol(optarg, &null_check, 10);
          if (*null_check) {
            errx(1, "ERROR: Args %s is not a valid integer.", optarg);
          }
        } else {
          printf("ERROR: no argument specified for musl_cpu\n");
          Usage();
        }
        break;
      /*case 'x':
        if (*optarg) {
          opts.xmlpath = optarg;
        } else {
          printf("ERROR: no argument specified for musl_xml\n");
          Usage();
        }
        break;*/
      case 'a':
        if (*optarg) {
          opts.extra_benchmarks.push_back(optarg);
        } else {
          printf("ERROR: no argument specified for musl_extra\n");
          Usage();
        }
        break;
      case 'i':
        if (*optarg){
          char* null_check;
          opts.iter_num = strtol(optarg, &null_check, 10);
          if (*null_check != '\0' or opts.iter_num < 0) {
            errx(1, "ERROR: Args %s is not a valid number of iterations.", optarg);
          }
        } else {
          printf("ERROR: no argument specified for musl_iterations\n");
          Usage();
        }
        break;
      case 'h':
        Usage();
        break;
      case '?':
        break;
      default:
        exit(1);
    }
  }
  return opts;
}

bool LockOnCPU(int cpu_num) {
  cpu_set_t cpuset;

  CPU_ZERO(&cpuset);
  CPU_SET(cpu_num, &cpuset);
  if (sched_setaffinity(0, sizeof(cpuset), &cpuset) != 0) {
    if (errno == EINVAL) {
      printf("Invalid cpu %d\n", cpu_num);
    } else {
      perror("sched_setaffinity set to cpu failed:");
    }
    return false;
  }

  return true;
}

void LockAndRun(benchmark::State& state, benchmark_func func_to_bench, int cpu_num) {
  if (cpu_num >= 0) {
    LockOnCPU(cpu_num);
  }

  reinterpret_cast<void(*) (benchmark::State&)>(func_to_bench)(state);
}

static constexpr char onebufManualStr[] = "ONEBUF_MANUAL_ALIGN_";
static constexpr char twobufManualStr[] = "TWOBUF_MANUAL_ALIGN1_";

static bool ParseOnebufManualStr(std::string& arg, args_vector* argsVector) {
  // The format of this is:
  //   ONEBUF_MANUAL_ALIGN_XX_SIZE_YY
  // Where:
  //   XX is the alignment
  //   YY is the size
  // The YY size can be either a number or a string representing the pre-defined
  // sets of values:
  //   SMALL (for values between 1 and 256)
  //   MEDIUM (for values between 512 and 128KB)
  //   LARGE (for values between 256KB and 2048KB)
  int64_t align;
  int64_t iSize;
  char strSizes[32] = { 0 };
  int ret;

  ret = sscanf(arg.c_str(), "ONEBUF_MANUAL_ALIGN_%" SCNd64 "_SIZE_%" SCNd64, &align, &iSize);
  if (ret == 1) {
    ret = sscanf(arg.c_str(), "ONEBUF_MANUAL_ALIGN_%" SCNd64 "_SIZE_"
                              "%" STRCONNECT(sizeof(strSizes)-1) "s", &align, strSizes);
  }
  if (ret != 2) {
    return false;
  }

  // Verify the alignment is powers of 2.
  if (align != 0 && (align & (align - 1)) != 0) {
    return false;
  }

  auto it = argsSize.find(strSizes);
  if (it == argsSize.cend()) {
    argsVector->push_back({iSize, align});
  } else {
    for (auto size : it->second) {
      argsVector->push_back({size, align});
    }
  }
  return true;
}

static bool ParseTwobufManualStr(std::string& arg, args_vector* argsVector) {
  // The format of this is:
  //   TWOBUF_MANUAL_ALIGN1_XX_ALIGN2_YY_SIZE_ZZ
  // Where:
  //   XX is the alignment of the first argument
  //   YY is the alignment of the second argument
  //   ZZ is the size
  // The ZZ size can be either a number or a string representing the pre-defined
  // sets of values:
  //   SMALL (for values between 1 and 256)
  //   MEDIUM (for values between 512 and 128KB)
  //   LARGE (for values between 256KB and 2048KB)
  int64_t align1;
  int64_t align2;
  int64_t iSize;
  char strSizes[32] = { 0 };
  int ret;

  ret = sscanf(arg.c_str(), "TWOBUF_MANUAL_ALIGN1_%" SCNd64 "_ALIGN2_%" SCNd64 "_SIZE_%" SCNd64,
                            &align1, &align2, &iSize);
  if (ret == 2) {
    ret = sscanf(arg.c_str(), "TWOBUF_MANUAL_ALIGN1_%" SCNd64 "_ALIGN2_%" SCNd64 "_SIZE_"
                               "%" STRCONNECT(sizeof(strSizes)-1) "s", &align1, &align2, strSizes);
  }
  if (ret != 3) {
    return false;
  }

  // Verify the alignments are powers of 2.
  if ((align1 != 0 && (align1 & (align1 - 1)) != 0)
      || (align2 != 0 && (align2 & (align2 - 1)) != 0)) {
    return false;
  }

  auto it = argsSize.find(strSizes);
  if (it == argsSize.cend()) {
    argsVector->push_back({iSize, align1, align2});
  } else {
    for (auto size : it->second) {
      argsVector->push_back({size, align1, align2});
    }
  }
  return true;
}

args_vector* ResolveArgs(args_vector* argsVector, std::string args,
                           std::map<std::string, args_vector>& argsAbbreviation) {
  // args is either a space-separated list of ints, a macro name, or
  // special free form macro.
  if (argsAbbreviation.count(args)) {
    return &argsAbbreviation[args];
  }
  // Check for free form macro.
  if (args.find_first_of(onebufManualStr) == 0) {
    if (!ParseOnebufManualStr(args, argsVector)) {
      errx(1, "ERROR: Bad format of macro %s, should be ONEBUF_MANUAL_ALIGN_XX_SIZE_YY", args.c_str());
    }
    return argsVector;
  } else if (args.find_first_of(twobufManualStr) == 0) {
    if (!ParseTwobufManualStr(args, argsVector)) {
      errx(1, "ERROR: Bad format of macro %s, should be TWOBUF_MANUAL_ALIGN1_XX_ALIGNE2_YY_SIZE_ZZ", args.c_str());
    }
    return argsVector;
  }

  argsVector->push_back(std::vector<int64_t>());
  std::stringstream sstream(args);
  std::string argstr;
  while (sstream >> argstr) {
    char* null_check;
    int converted = static_cast<int>(strtol(argstr.c_str(), &null_check, 10));
    if (*null_check) {
      errx(1, "ERROR: Args str %s contains an invalid macro or int.", args.c_str());
    }
    (*argsVector)[0].push_back(converted);
  }
  return argsVector;
}

void RegisterOriginalBenchmarks(bench_opts_t opts, const std::string& fn_name, args_vector* run_args) {
  if (g_str_func_map.find(fn_name) == g_str_func_map.end()) {
    errx(1, "ERROR: No benchmark for function %s", fn_name.c_str());
  }

  benchmark_func func = g_str_func_map.at(fn_name).first;
  for (const std::vector<int64_t>& args : (*run_args)) {
    auto registration = benchmark::RegisterBenchmark(fn_name.c_str(), LockAndRun, func, opts.cpu_num)->Args(args);
    if (opts.iter_num > 0) {
      registration->Iterations(opts.iter_num);
    }
  }
}

std::string Trim(std::string s) {
  if(s.empty()) {
      return s;
  }
  s.erase(0, s.find_first_not_of(" "));
  s.erase(s.find_last_not_of(" ") + 1);
  return s;
}

void RegisterExtraBenchmarks(bench_opts_t cmdline_opts, std::map<std::string, args_vector>& argsAbbreviation) {
  // Register any of the extra benchmarks that were specified in the options.
  args_vector arg_vector;
  args_vector* run_args = &arg_vector;
  std::string fn_name;
  std::string cmd_args;
  for (const std::string& extra_fn : cmdline_opts.extra_benchmarks) {
    Trim(extra_fn);
    size_t first_space_pos = extra_fn.find(' ');
    fn_name = extra_fn.substr(0, first_space_pos);
    if (first_space_pos != std::string::npos) {
      cmd_args = extra_fn.substr(extra_fn.find(' ') + 1);
    } else {
      cmd_args = "";
    }
    run_args = ResolveArgs(run_args, cmd_args, argsAbbreviation);
    RegisterOriginalBenchmarks(cmdline_opts, fn_name, run_args);

    run_args = &arg_vector;
    arg_vector.clear();
  }
}

int RegisterXmlBenchmarks(bench_opts_t cmdline_opts, std::map<std::string, args_vector>& argsAbbreviation) {
  // Structure of the XML file:
  // - Element "fn"           Function to benchmark.
  // - - Element "iterations" Number of iterations to run. Leaving this blank uses
  //                          benchmarks' convergence heuristics.
  // - - Element "cpu"        CPU to isolate to, if any.
  // - - Element "args"       Whitespace-separated list of per-function integer arguments, or
  //                          one of the macros defined in util.h.
#if 0
  tinyxml2::XMLDocument doc;
  if (doc.LoadFile(cmdline_opts.xmlpath.c_str()) != tinyxml2::XML_SUCCESS) {
    doc.PrintError();
    return doc.ErrorID();
  }

  // Read and register the functions.
  tinyxml2::XMLNode* fn = doc.FirstChildElement("fn");
  while (fn) {
    if (fn == fn->ToComment()) {
      // Skip comments.
      fn = fn->NextSibling();
      continue;
    }

    auto fn_elem = fn->FirstChildElement("name");
    if (!fn_elem) {
      errx(1, "ERROR: Malformed XML entry: missing name element.");
    }
    std::string fn_name = fn_elem->GetText();
    if (fn_name.empty()) {
      errx(1, "ERROR: Malformed XML entry: error parsing name text.");
    }
    auto* xml_args = fn->FirstChildElement("args");
    args_vector arg_vector;
    args_vector* run_args = ResolveArgs(&arg_vector,
                                          xml_args ? Trim(xml_args->GetText()) : "",
                                          argsAbbreviation);

    // XML values for CPU and iterations take precedence over those passed in via CLI.
    bench_opts_t xml_opts{};
    auto* num_iterations_elem = fn->FirstChildElement("iterations");
    if (num_iterations_elem) {
      int temp;
      num_iterations_elem->QueryIntText(&temp);
      xml_opts.iter_num = temp;
    }
    auto* cpu_to_lock_elem = fn->FirstChildElement("cpu");
    if (cpu_to_lock_elem) {
      int temp;
      cpu_to_lock_elem->QueryIntText(&temp);
      xml_opts.cpu_num = temp;
    }

    RegisterOriginalBenchmarks(xml_opts, cmdline_opts, fn_name, run_args);

    fn = fn->NextSibling();
  }
#endif
  return 0;
}

static void SetArgs(const std::vector<int>& sizes, args_vector* args) {
  for (int size : sizes) {
    args->push_back({size});
  }
}

static void SetArgs(const std::vector<int>& sizes, int align, args_vector* args) {
  for (int size : sizes) {
    args->push_back({size, align});
  }
}


static void SetArgs(const std::vector<int>& sizes, int align1, int align2, args_vector* args) {
  for (int size : sizes) {
    args->push_back({size, align1, align2});
  }
}

static args_vector GetArgs(const std::vector<int>& sizes) {
  args_vector args;
  SetArgs(sizes, &args);
  return args;
}

static args_vector GetArgs(const std::vector<int>& sizes, int align) {
  args_vector args;
  SetArgs(sizes, align, &args);
  return args;
}

static args_vector GetArgs(const std::vector<int>& sizes, int align1, int align2) {
  args_vector args;
  SetArgs(sizes, align1, align2, &args);
  return args;
}

std::map<std::string, args_vector> GetAbbreviation() {
  std::vector<int> allArgs(smallArgs);
  allArgs.insert(allArgs.end(), mediumArgs.begin(), mediumArgs.end());
  allArgs.insert(allArgs.end(), largeArgs.begin(), largeArgs.end());

  std::map<std::string, args_vector> argsAbbreviation {
    {"COMMON_ARGS", GetArgs(commonArgs)},
    {"SMALL_ARGS", GetArgs(smallArgs)},
    {"MEDIUM_ARGS", GetArgs(mediumArgs)},
    {"LARGE_ARGS", GetArgs(largeArgs)},
    {"ALL_ARGS", GetArgs(allArgs)},

    {"ALIGNED_ONEBUF", GetArgs(commonArgs, 0)},
    {"ALIGNED_ONEBUF_SMALL", GetArgs(smallArgs, 0)},
    {"ALIGNED_ONEBUF_MEDIUM", GetArgs(mediumArgs, 0)},
    {"ALIGNED_ONEBUF_LARGE", GetArgs(largeArgs, 0)},
    {"ALIGNED_ONEBUF_ALL", GetArgs(allArgs, 0)},

    {"ALIGNED_TWOBUF", GetArgs(commonArgs, 0, 0)},
    {"ALIGNED_TWOBUF_SMALL", GetArgs(smallArgs, 0, 0)},
    {"ALIGNED_TWOBUF_MEDIUM", GetArgs(mediumArgs, 0, 0)},
    {"ALIGNED_TWOBUF_LARGE", GetArgs(largeArgs, 0, 0)},
    {"ALIGNED_TWOBUF_ALL", GetArgs(allArgs, 0, 0)},

    {"MATH_COMMON", args_vector{ {0}, {1}, {2}, {3}, {4}, {5} }},
    {"NANO_SLEEP", args_vector{ {0}, {10000},{100000000} }},
    {"MALLOC_FREE", args_vector{ {1}, {10},{100},{1000},{10000} }},
    {"PTHREAD_KEY",args_vector{{19},{38},{39}}},
  };

  args_vector argsOnebuf;
  args_vector argsTwobuf;
  for (int size : allArgs) {
    argsOnebuf.push_back({size, 0});
    argsTwobuf.push_back({size, 0, 0});
    // Skip alignments on zero sizes.
    if (size == 0) {
      continue;
    }
    for (int align1 = 1; align1 <= 32; align1 <<= 1) {
      argsOnebuf.push_back({size, align1});
      for (int align2 = 1; align2 <= 32; align2 <<= 1) {
        argsTwobuf.push_back({size, align1, align2});
      }
    }
  }
  argsAbbreviation.emplace("MANY_ALIGNED_ONEBUF", argsOnebuf);
  argsAbbreviation.emplace("MANY_ALIGNED_TWOBUF", argsTwobuf);

  return argsAbbreviation;
}
/*
static bool FileExists(const std::string& file) {
  struct stat st;
  return stat(file.c_str(), &st) != -1 && S_ISREG(st.st_mode);
}*/

void RegisterAllBenchmarks(const bench_opts_t& opts, std::map<std::string, args_vector>& argsAbbreviation) {
  for (auto& entry : g_str_func_map) {
    auto& function_info = entry.second;
    args_vector arg_vector;
    args_vector* run_args = ResolveArgs(&arg_vector, function_info.second, argsAbbreviation);
    RegisterOriginalBenchmarks(opts, entry.first, run_args);
  }
}

int main(int argc, char** argv) {
  std::map<std::string, args_vector> argsAbbreviation = GetAbbreviation();
  bench_opts_t opts = ParseOpts(argc, argv);
  std::vector<char*> argv_for_origin(argc);
  SanitizeOpts(argc, argv, &argv_for_origin);

//  if (opts.xmlpath.empty()) {
    // Don't add the default xml file if a user is specifying the tests to run.
  if (opts.extra_benchmarks.empty()) {
    RegisterAllBenchmarks(opts, argsAbbreviation);
  } else {
    RegisterExtraBenchmarks(opts, argsAbbreviation);
  }
  /*} else if (!FileExists(opts.xmlpath)) {
    // See if this is a file in the suites directory.
    std::string file(android::base::GetExecutableDirectory() + "/suites/" + opts.xmlpath);
    if (opts.xmlpath[0] == '/' || !FileExists(file)) {
      printf("Cannot find xml file %s: does not exist or is not a file.\n", opts.xmlpath.c_str());
      return 1;
    }
    opts.xmlpath = file;
  }

  if (!opts.xmlpath.empty()) {
    if (int err = RegisterXmlBenchmarks(opts, argsAbbreviation)) {
      return err;
    }
  }*/
  

  // Set the thread priority to the maximum.
  if (setpriority(PRIO_PROCESS, 0, -20)) {
    perror("Failed to raise priority of process. Are you root?\n");
  }

  int argc_for_origin = argv_for_origin.size();
  benchmark::Initialize(&argc_for_origin, argv_for_origin.data());
  benchmark::RunSpecifiedBenchmarks();
}
